0398. 随机数索引【中等】

• 1. 📝 题目描述
• 2. 🎯 s.1 - 蓄水池抽样

1. 📝 题目描述

• leetcode

给你一个可能含有 重复元素 的整数数组 nums，请你随机输出给定的目标数字 target 的索引。你可以假设给定的数字一定存在于数组中。

实现 Solution 类：

• Solution(int[] nums) 用数组 nums 初始化对象。
• int pick(int target) 从 nums 中选出一个满足 nums[i] == target 的随机索引 i。如果存在多个有效的索引，则每个索引的返回概率应当相等。

示例：

输入
["Solution", "pick", "pick", "pick"]
[[[1, 2, 3, 3, 3]], [3], [1], [3]]
输出
[null, 4, 0, 2]

解释
Solution solution = new Solution([1, 2, 3, 3, 3]);
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
solution.pick(1); // 返回 0。因为只有 nums[0] 等于 1。
solution.pick(3); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。

---

提示：

• 1 <= nums.length <= 2 * 10^4
• -2^31 <= nums[i] <= 2^31 - 1
• target 是 nums 中的一个整数
• 最多调用 pick 函数 10^4 次

2. 🎯 s.1 - 蓄水池抽样

c

typedef struct {
    int* nums;
    int size;
} Solution;

Solution* solutionCreate(int* nums, int numsSize) {
    Solution* obj = (Solution*)malloc(sizeof(Solution));
    obj->nums = nums;
    obj->size = numsSize;
    return obj;
}

int solutionPick(Solution* obj, int target) {
    int res = 0, count = 0;
    for (int i = 0; i < obj->size; i++) {
        if (obj->nums[i] == target) {
            count++;
            if (rand() % count == 0) res = i;
        }
    }
    return res;
}

void solutionFree(Solution* obj) { free(obj); }

js

/**
 * @param {number[]} nums
 */
var Solution = function (nums) {
  this.nums = nums
}

/**
 * @param {number} target
 * @return {number}
 */
Solution.prototype.pick = function (target) {
  let res = 0,
    count = 0
  for (let i = 0; i < this.nums.length; i++) {
    if (this.nums[i] === target) {
      count++
      if (Math.random() < 1 / count) res = i
    }
  }
  return res
}

py

class Solution:
    def __init__(self, nums: List[int]):
        self.nums = nums

    def pick(self, target: int) -> int:
        res, count = 0, 0
        for i, num in enumerate(self.nums):
            if num == target:
                count += 1
                if random.randint(1, count) == 1:
                    res = i
        return res

• 时间复杂度：初始化 $O(1)$，pick $O(n)$
• 空间复杂度：$O(1)$（不计输入数组）

算法思路：

• 蓄水池抽样：遍历时，第 $k$ 个等于 target 的元素以 $\frac{1}{k}$ 的概率被选中
• 保证每个符合条件的索引被等概率返回